Skip to content

[Bugfix][Model Runner V2][Spec Decode] Decouple the draft's gumbel noise stream from the target's - #54282

Merged
WoosukKwon merged 1 commit into
vllm-project:mainfrom
TheEpicDolphin:mrv2-draft-gumbel-decouple
Aug 29, 2026
Merged

WoosukKwon merged 1 commit into
vllm-project:mainfrom
TheEpicDolphin:mrv2-draft-gumbel-decouple

Conversation

@TheEpicDolphin

@TheEpicDolphin TheEpicDolphin commented Aug 28, 2026 •

Copy link
Copy Markdown
Collaborator

Context

Picking up the work from #47386. During target verification, the Gumbel noise used to re-sample a rejected draft token is drawn from the same Philox offset as the noise that produced that draft token (the same seed, pos), so byte-identical noise. Conditioned on the proposal winning the argmax, the other tokens' Gumbels are truncated below that max, most tightly for the tokens the draft ranked highest, but lost (high-q loser tokens generally have smaller noise, because otherwise they would have won), so the residual under-weights exactly those tokens and the output distribution is no longer the target's.

This affects draft_sample_method="probabilistic" only. Under the default greedy, draft_logits is None, the draft never calls gumbel_sample, and there is no shared noise vector.

This PR

Adds a salt (1 << 30) to the offset of the draft's Philox stream, decoupling the draft's Gumbel noise from the target's. Verification is a probability-ratio test against the cached draft distribution rather than a Gumbel coupling, so nothing depends on the two streams matching.

Also refactors how the speculators compute the position used as the sampling-stream index, so each site states its own convention instead of relying on an implicit +1 inside sample_draft. This incidentally fixes a second latent issue: the sampling position was previously derived from a buffer clamped at max_model_len, so near the context limit two consecutive draft steps could reuse one noise vector. It is now tracked unclamped.

Tests

gdelfin@gb200-rack1-18:~/vllm$ python -m pytest tests/v1/spec_decode/test_rejection_sampler_utils.py::test_gumbel_drafted_rejection_sample_is_unbiased -q
Running 2 items in this shard
..                                                                                                                                                                                                                                                                                                                          [100%]
2 passed in 2.96s

gdelfin@gb200-rack1-18:~/vllm$ python -m pytest tests/v1/worker/test_gpu_gumbel_sample.py::test_drafting_uses_a_separate_noise_stream -q
Running 1 items in this shard
.                                                                                                                                                                                                                                                                                                                           [100%]
1 passed in 3.92s

Benchmarks

vllm bench serve --dataset-name speed_bench --speed-bench-dataset-subset throughput_2k --speed-bench-output-len 2048
— 2K in / 2K out, 512 prompts, 64 warmups, concurrency 64, temperature=1.0 top_p=0.95 top_k=20.
All runs use draft_sample_method: probabilistic (the salt is a no-op under the default greedy,
where the draft never calls gumbel_sample). 4× GB200, TP4 except Gemma 4 (TP1).
Baseline de9250ac9e → this PR 1529dd5363.

Model N Output tok/s Acceptance length Acceptance rate % Median TTFT (ms) Median TPOT (ms) Median ITL (ms)
DeepSeek V4 Flash + MTP 3 5444 → 5438 (-0.11%) 2.589 → 2.587 (-0.09%) 52.96 → 52.89 (-0.14%) 221.0 → 226.1 (+2.30%) 11.64 → 11.60 (-0.29%) 20.11 → 20.11 (+0.01%)
DeepSeek V4 Flash + DSpark 5 5825 → 5942 (+2.00%) 3.514 → 3.500 (-0.39%) 50.28 → 50.00 (-0.55%) 237.1 → 225.9 (-4.74%) 10.83 → 10.52 (-2.82%) 24.05 → 23.99 (-0.23%)
MiMo v2.5 Pro + DFlash 8 4111 → 4086 (-0.62%) 3.794 → 3.771 (-0.61%) 34.93 → 34.64 (-0.83%) 210.2 → 206.7 (-1.68%) 14.44 → 14.43 (-0.05%) 44.14 → 44.15 (+0.03%)
Inkling NVFP4 + MTP 8 3813 → 3864 (+1.33%) 3.381 → 3.386 (+0.15%) 29.76 → 29.82 (+0.22%) 205.4 → 212.8 (+3.59%) 16.22 → 16.07 (-0.92%) 47.40 → 47.31 (-0.18%)
Gemma 4 E2B + MTP 3 10609 → 10508 (-0.96%) 2.623 → 2.612 (-0.39%) 54.09 → 53.75 (-0.64%) 92.5 → 94.5 (+2.18%) 5.48 → 5.67 (+3.58%) 6.80 → 6.69 (-1.54%)

Per-position acceptance rate (%), baseline → this PR:

Model pos 0 pos 1 pos 2 pos 3 pos 4 pos 5 pos 6 pos 7
DeepSeek V4 Flash + MTP 85.2 → 85.1 53.1 → 53.1 20.6 → 20.5 — — — — —
DeepSeek V4 Flash + DSpark 83.6 → 83.5 64.6 → 64.4 47.3 → 46.9 33.1 → 32.8 22.8 → 22.5 — — —
MiMo v2.5 Pro + DFlash 76.5 → 76.3 56.1 → 55.9 42.2 → 41.9 32.5 → 32.3 25.5 → 25.1 19.9 → 19.5 15.4 → 15.1 11.4 → 11.0
Inkling NVFP4 + MTP 68.5 → 68.6 48.0 → 48.1 34.5 → 34.6 26.0 → 26.1 20.2 → 20.3 16.3 → 16.3 13.4 → 13.4 11.2 → 11.2
Gemma 4 E2B + MTP 70.8 → 70.5 51.8 → 51.4 39.7 → 39.3 — — — — —

No regressions. Throughput deltas straddle zero (-0.96% to +2.00%) and acceptance length is
within ±0.61% with both signs present — consistent with run-to-run variance rather than a
systematic effect. Per-position rates show no depth-dependent degradation: MiMo's 8-deep profile
moves uniformly by -0.2 to -0.4 pp with no trend across depth.

Acceptance is not expected to move: the acceptance test is p(x) > u·q(x), whose expectation
depends only on x ~ q, which holds in both arms. An isolated probe with fixed p/q over 2M
independent trials put acceptance at 0.1068 (baseline) vs 0.1071 (this PR).

…ise stream from the target's

Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai>
@mergify mergify Bot added speculative-decoding dflash mrv2 Model Runner V2 specific bug Something isn't working labels Aug 28, 2026
@TheEpicDolphin
TheEpicDolphin marked this pull request as ready for review August 29, 2026 01:16

@claude claude Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

@TheEpicDolphin TheEpicDolphin added the ready ONLY add when PR is ready to merge/full CI is needed label Aug 29, 2026
@TheEpicDolphin

Copy link
Copy Markdown
Collaborator Author

/ci run

@github-actions

Copy link
Copy Markdown

✅ Triggered Buildkite CI #86086 for commit 1529dd536348.

@WoosukKwon WoosukKwon left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you!

@WoosukKwon
WoosukKwon merged commit fe755c8 into vllm-project:main Aug 29, 2026
124 checks passed
@github-project-automation github-project-automation Bot moved this from Backlog to Done in Sprint - DFlash Aug 29, 2026
randomvariable added a commit to randomvariable/vllm that referenced this pull request Aug 29, 2026
Upstream's draft/target noise-stream decoupling (vllm-project#54282) put the fork's
own mechanism inside the Gumbel kernel: `_DRAFT_NOISE_SALT` is the same
`1 << 30` offset, applied under `IS_DRAFTING`, keyed off a persistent
`sample_src_positions` buffer that the draft-input kernel advances
device-side regardless of `ADVANCE_DRAFT_POSITIONS`.

Keeping `draft_gumbel_pos` on top of that would salt the draw twice, so
delete it and let the one central sampler carry the intent:

* `_sample_probabilistic_draft` forwards `sample_src_positions`
  untouched with `is_drafting=True`, and still forwards `active_rows`
  for the variable-width autoregressive path.
* `DraftModelSpeculator.sample_draft` and the DSpark/DFlash overrides
  key by the position before the sampled token, P-1, matching upstream.
* The autoregressive override's parameter takes the base class's name.

Tests follow the mechanism rather than the retired helper: the sampler
test asserts pass-through positions plus `is_drafting`, and the
rejection-sampling distribution test keys the draft draw the same way
production does.

Verified: `tools/ci/static_gate.py` reports no undeclared envs
references and no call-arg/attr-defined drift, at the same 792/139 mypy
debt count as the pre-rebase tip; the ruff F821/F811 set over `vllm`
and `tests` is identical to that tip; 4220 files parse.

Co-authored-by: OMP Agent <noreply@omp.local>
Signed-off-by: Naadir Jeewa <naadir@randomvariable.co.uk>
@TheEpicDolphin
TheEpicDolphin deleted the mrv2-draft-gumbel-decouple branch August 31, 2026 20:08
am-cohere pushed a commit to am-cohere/vllm that referenced this pull request Sep 1, 2026
…ise stream from the target's (vllm-project#54282)

Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai>
Leoyzen pushed a commit to Leoyzen/vllm that referenced this pull request Sep 1, 2026
…ise stream from the target's (vllm-project#54282)

Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai>
mylibrar pushed a commit to tanyuqian/vllm that referenced this pull request Sep 3, 2026
…ise stream from the target's (vllm-project#54282)

Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai>
D-G-Dimitrov pushed a commit to D-G-Dimitrov/vllm that referenced this pull request Sep 5, 2026
…ise stream from the target's (vllm-project#54282)

Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai>
(cherry picked from commit fe755c8)
voipmonitor pushed a commit to local-inference-lab/vllm that referenced this pull request Sep 11, 2026
…roject#54282

Adapt upstream vllm-project/vllm commit fe755c8 for Jovian Judgement. Preserve compact MTP prefill outputs, MRoPE, per-token draft-logit caching, and local DFlash2/DSpark call sites. Extend BF16 proposal and fixed/clamped-position tests.

Original fix by Giancarlo Delfin; JJ adaptation and additional regression coverage by MadeBy561 with Codex assistance. Prior R25 GPU results are not fresh-head qualification.

Co-authored-by: MadeBy561 <madeby561@gmail.com>
Co-authored-by: Codex <noreply@openai.com>
Signed-off-by: MadeBy561 <madeby561@gmail.com>
voipmonitor added a commit to local-inference-lab/vllm that referenced this pull request Sep 11, 2026
…upstream vllm-project#54282

Preserve the reviewed source head 44e6766 and its contributor history.
The first parent records the ordered serving-source composition.
Whole-tree equality and installed-artifact verification are publication gates.

Review: #653
Assisted-by: OpenAI Codex
Signed-off-by: Martin Vit <martin@voipmonitor.org>
Terrydaktal added a commit to Terrydaktal/vllm-radiance that referenced this pull request Sep 14, 2026
The DFlash2 selector used the same Philox noise stream as the target's
replacement draw after rejection. This conditions replacements on the rejected
proposal and changes the intended output distribution. Adapt the independent
draft stream from vllm-project/vllm#54282 (fe755c88995ad468882517b6c4bdd60138d46a3a)
to the selector's direct Gumbel call. Offset only its local RNG index, preserve
model/cache positions and greedy behavior, and install the guarded idempotent
patch in the release, development and patch-image build paths.

The pinned libr4d GDN scan's midpoint factorization clamps growing exponents
at 80. Large decay spans can erase valid contributions and final recurrent
state without producing nonfinite values. Extend the existing libr4d build
patch with a bounded FP32 recurrence for affected sequence/head pairs, using
a conservative span threshold of 128. Preserve ordinary fast-scan outputs,
the entry-point ABI, caller stream, tensor layouts and decode path. Propagate
the original launch error before launching the correction. No host readback,
synchronization or additional temporary tensor is introduced.

Add public, checkpoint-free GPU regressions for native selector/rejection
sampling and an independent FP64 GDN recurrence. Include ordinary and extreme
decays, unequal sequences, partial chunks, mixed heads, large values, both
48/16 and 24/8 head layouts, and graph replays with changed inputs. Record
compact numeric qualification evidence and document reproduction, provenance,
the need to rebuild old computed state, and remaining performance limits.
Link the correction documentation from the README.

Validation on R9700/gfx1201 with the pinned Radiance 1.0.16 image:
- Patch applies to pinned libr4d source; both scan units compile with its
  production compiler flags; sampler installation and repeat installation pass.
- At three positions with 200,000 draws per arm, shared-noise probability error
  is 0.01781-0.01953 versus 0.000965-0.001180 after correction. Target-only and
  greedy controls pass.
- Original GDN reproduces 46.4% and 82.8% output error and 100% state error
  on two extreme-decay cases. Both corrected layouts pass all 12 cases and
  three graph replays, with output error below 0.345% and state error below
  0.252% against FP64.

This commit is limited to the two numerical defects and their validation.
End-to-end throughput, distributed TP2 and repetition-rate improvements are
not claimed. The GDN correction adds prefill work that still needs workload
performance qualification. No loop guard, history recovery or serving-policy
change is included.
Terrydaktal added a commit to Terrydaktal/vllm-coherence that referenced this pull request Sep 18, 2026
…nt random streams

Backport the independently qualified upstream random-stream salt so a rejected draft does not reuse its Gumbel draw for the target replacement. Preserve the declared sampling configuration rather than hiding the bias with different decoding settings.

Authenticate the affected source preimage and include CPU regressions for repeatable stream separation. Credit the upstream fix rather than treating it as a new discovery.

Provenance
Upstream PR: vllm-project/vllm#54282 (author: TheEpicDolphin).
This is an adapted backport for the pinned vLLM API.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working dflash mrv2 Model Runner V2 specific ready ONLY add when PR is ready to merge/full CI is needed speculative-decoding

Projects

Status: Done

Development

Successfully merging this pull request may close these issues.

2 participants